{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# default_exp export"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# export: nb2py\n",
    "\n",
    "> Code that allows you to export a notebook (.ipynb) as a python script( .py) to a target folder."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "nb2py will allow you to convert the notebook (.ipynb) where the function is executed to a python script. \n",
    "\n",
    "The conversion applies these rules: \n",
    "\n",
    "* The notebook will be automatically saved when the function is executed.\n",
    "* Only code cells will be converted (not markdown cells).\n",
    "* A header will be added to indicate the script has been automatically generated. It also indicates where the original ipynb is.\n",
    "* Cells with a #hide flag won't be converted. Flag variants like # hide, #Hide, #HIDE, ... are also acceptable.\n",
    "* Empty cells and unnecessary empty lines within cells will be removed.\n",
    "* By default the script will be created with the same name and in the same folder of the original notebook. But you can pass a dir folder and a different name if you wish.\n",
    "* If a script with the same name already exists, it will be overwriten."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "from tsai.imports import *\n",
    "from fastcore.script import *\n",
    "import runpy\n",
    "import re\n",
    "import nbformat\n",
    "import requests\n",
    "from urllib.parse import unquote\n",
    "import shutil"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This code is required to identify flags in the notebook. We are looking for #hide flags."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "def maybe_mount_gdrive():\n",
    "    from pathlib import Path\n",
    "    from google.colab.drive import mount\n",
    "    if not Path(\"/content/drive\").exists(): mount(\"/content/drive\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "\"\"\"\n",
    "Code in this cell is based on the nbdev library: \n",
    "https://github.com/fastai/nbdev/blob/master/nbdev/export.py\n",
    "\"\"\"\n",
    "\n",
    "def _mk_flag_re(body, n_params, comment):\n",
    "    \"Compiles a regex for finding nbdev flags\"\n",
    "    assert body!=True, 'magics no longer supported'\n",
    "    prefix = r\"\\s*\\#\\s*\"\n",
    "    param_group = \"\"\n",
    "    if n_params == -1: param_group = r\"[ \\t]+(.+)\"\n",
    "    if n_params == 1: param_group = r\"[ \\t]+(\\S+)\"\n",
    "    if n_params == (0,1): param_group = r\"(?:[ \\t]+(\\S+))?\"\n",
    "    return re.compile(rf\"\"\"\n",
    "# {comment}:\n",
    "^            # beginning of line (since re.MULTILINE is passed)\n",
    "{prefix}\n",
    "{body}\n",
    "{param_group}\n",
    "[ \\t]*       # any number of spaces and/or tabs\n",
    "$            # end of line (since re.MULTILINE is passed)\n",
    "\"\"\", re.MULTILINE | re.VERBOSE)\n",
    "\n",
    "_re_hide = _mk_flag_re(\"hide?\", 0,\n",
    "    \"Matches any line with #hide without any module name\")\n",
    "\n",
    "def _get_unhidden_cells(cells):\n",
    "    result = []\n",
    "    for i,cell in enumerate(cells):\n",
    "        if cell['cell_type'] == 'code': \n",
    "            if not _re_hide.findall(cell['source'].lower()) and cell['source'] != '': result.append(i)\n",
    "    return result\n",
    "    \n",
    "def _read_nb(fname):\n",
    "    \"Read the notebook in `fname`.\"\n",
    "    with open(Path(fname),'r', encoding='utf8') as f: return nbformat.reads(f.read(), as_version=4)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This code automatically gets the name of the notebook. It's been tested to work on Jupyter notebooks, Jupyter Lab and Google Colab."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "\"\"\"Code in this cell is a modified version of this repo: \n",
    "https://github.com/msm1089/ipynbname\n",
    "# Copyright (c) 2020 Mark McPherson. \"\"\"\n",
    "\n",
    "# MIT License\n",
    "\n",
    "# Copyright (c) 2020 Mark McPherson\n",
    "\n",
    "# Permission is hereby granted, free of charge, to any person obtaining a copy\n",
    "# of this software and associated documentation files (the \"Software\"), to deal\n",
    "# in the Software without restriction, including without limitation the rights\n",
    "# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell\n",
    "# copies of the Software, and to permit persons to whom the Software is\n",
    "# furnished to do so, subject to the following conditions:\n",
    "\n",
    "# The above copyright notice and this permission notice shall be included in all\n",
    "# copies or substantial portions of the Software.\n",
    "\n",
    "# THE SOFTWARE IS PROVIDED \"AS IS\", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR\n",
    "# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,\n",
    "# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE\n",
    "# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER\n",
    "# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,\n",
    "# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE\n",
    "# SOFTWARE.\n",
    "\n",
    "import json\n",
    "import urllib.error\n",
    "import urllib.request\n",
    "from itertools import chain\n",
    "from pathlib import Path, PurePath\n",
    "from typing import Generator, Tuple, Union\n",
    "import ipykernel\n",
    "from jupyter_core.paths import jupyter_runtime_dir\n",
    "from traitlets.config import MultipleInstanceError\n",
    "\n",
    "_FILE_ERROR = \"Can't identify the notebook {}.\"\n",
    "_CONN_ERROR = \"Unable to access server;\\n\" \\\n",
    "           + \"ipynbname requires either no security or token based security.\"\n",
    "\n",
    "def _list_maybe_running_servers(runtime_dir=None) -> Generator[dict, None, None]:\n",
    "    \"\"\" Iterate over the server info files of running notebook servers.\n",
    "    \"\"\"\n",
    "    if runtime_dir is None:\n",
    "        runtime_dir = jupyter_runtime_dir()\n",
    "    runtime_dir = Path(runtime_dir)\n",
    "\n",
    "    if runtime_dir.is_dir():\n",
    "        for file_name in chain(\n",
    "            runtime_dir.glob('nbserver-*.json'),  # jupyter notebook (or lab 2)\n",
    "            runtime_dir.glob('jpserver-*.json'),  # jupyterlab 3\n",
    "        ):\n",
    "            yield json.loads(file_name.read_bytes())\n",
    "\n",
    "def _get_kernel_id() -> str:\n",
    "    \"\"\" Returns the kernel ID of the ipykernel.\n",
    "    \"\"\"\n",
    "    connection_file = Path(ipykernel.get_connection_file()).stem\n",
    "    kernel_id = connection_file.split('-', 1)[1]\n",
    "    return kernel_id\n",
    "\n",
    "def _get_sessions(srv):\n",
    "    \"\"\" Given a server, returns sessions, or HTTPError if access is denied.\n",
    "        NOTE: Works only when either there is no security or there is token\n",
    "        based security. An HTTPError is raised if unable to connect to a \n",
    "        server.\n",
    "    \"\"\"\n",
    "    try:\n",
    "        qry_str = \"\"\n",
    "        token = srv['token']\n",
    "        if token:\n",
    "            qry_str = f\"?token={token}\"\n",
    "        url = f\"{srv['url']}api/sessions{qry_str}\"\n",
    "        with urllib.request.urlopen(url) as req:\n",
    "            return json.load(req)\n",
    "    except Exception:\n",
    "        raise urllib.error.HTTPError(_CONN_ERROR)\n",
    "\n",
    "def _find_nb() -> Union[Tuple[dict, PurePath], Tuple[None, None]]:\n",
    "    try:\n",
    "        kernel_id = _get_kernel_id()\n",
    "    except (MultipleInstanceError, RuntimeError):\n",
    "        return None, None  # Could not determine\n",
    "    for srv in _list_maybe_running_servers():\n",
    "        try:\n",
    "            sessions = _get_sessions(srv)\n",
    "            for sess in sessions:\n",
    "                if sess['kernel']['id'] == kernel_id:\n",
    "                    return srv, PurePath(sess['notebook']['path'])\n",
    "        except Exception:\n",
    "            pass  # There may be stale entries in the runtime directory\n",
    "    return None, None\n",
    "\n",
    "def get_nb_name() -> str:\n",
    "    \"\"\" Returns the short name of the notebook w/o the .ipynb extension,\n",
    "        or raises a FileNotFoundError exception if it cannot be determined.\n",
    "    \"\"\"\n",
    "    try: \n",
    "        _, path = _find_nb()\n",
    "        if path:\n",
    "            return path.name\n",
    "        else:\n",
    "            return\n",
    "    except: \n",
    "        return\n",
    "\n",
    "def get_colab_nb_name():\n",
    "    d = requests.get('http://172.28.0.2:9000/api/sessions').json()[0]\n",
    "    fname = unquote(d['name'])\n",
    "    fid = unquote(d['path'].split('=')[1])\n",
    "    if 'https://github.com' in fid: fname = fid\n",
    "    else: fname = Path(f'drive/MyDrive/Colab Notebooks/{fname}')\n",
    "    return fname\n",
    "\n",
    "def get_nb_path() -> Path:\n",
    "    \"\"\" Returns the absolute path of the notebook,\n",
    "        or raises a FileNotFoundError exception if it cannot be determined.\n",
    "    \"\"\"\n",
    "    try: \n",
    "        if is_colab(): return get_colab_nb_name()\n",
    "        else: \n",
    "            srv, path = _find_nb()\n",
    "            if srv and path:\n",
    "                root_dir = Path(srv.get('root_dir') or srv['notebook_dir'])\n",
    "                return root_dir / path\n",
    "            else:\n",
    "                return\n",
    "    except: \n",
    "        return\n",
    "\n",
    "def nb_name_to_py(nb_name): \n",
    "    return str(nb_name).replace(\".ipynb\", \".py\")\n",
    "\n",
    "def get_script_path(nb_name=None):\n",
    "    if nb_name is None: nb_name = get_nb_path()\n",
    "    return nb_name_to_py(nb_name)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This code is used when trying to save a file to google drive. We first need to mount the drive."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "@call_parse\n",
    "def nb2py(nb:      Param(\"absolute or relative full path to the notebook you want to convert to a python script\", str)=None,\n",
    "          folder:  Param(\"absolute or relative path to folder of the script you will create. Defaults to current nb's directory\", str)=None,\n",
    "          name:    Param(\"name of the script you want to create. Defaults to current nb name .ipynb by .py\", str)=None,\n",
    "          save:    Param(\"saves the nb before converting it to a script\", store_false)=True,\n",
    "          run:     Param(\"import and run the script\", store_true)=False,\n",
    "          verbose: Param(\"controls verbosity\", store_false)=True,\n",
    "         ):\n",
    "    \"Converts a notebook to a python script in a predefined folder.\"\n",
    "    \n",
    "    # make sure drive is mounted when using Colab\n",
    "    if is_colab(): maybe_mount_drive()\n",
    "    \n",
    "    # nb path & name\n",
    "    if nb is not None:\n",
    "        nb_path = Path(nb)\n",
    "        nb_path = nb_path.parent/f\"{nb_path.stem}.ipynb\"\n",
    "    else:\n",
    "        try: \n",
    "            nb_path = get_nb_path()\n",
    "        except: \n",
    "            print(\"nb2py couldn't get the nb name. Pass it as an nb argument and rerun nb2py.\")\n",
    "            return\n",
    "    if nb_path is None: \n",
    "        print(\"nb2py couldn't get the nb name. Pass it as an nb argument and rerun nb2py.\")\n",
    "        return\n",
    "    nb_name = nb_path.name\n",
    "    assert os.path.isfile(nb_path), f\"nb2py couldn't find {nb_path}. Please, confirm the path is correct.\"\n",
    "    \n",
    "    # save nb: only those that are run from the notebook itself\n",
    "    if save and not is_colab() and nb is None: \n",
    "        try: save_nb(nb_name)\n",
    "        except: print(f\"nb2py couldn't save the nb automatically. It will used last saved at {to_local_time(os.path.getmtime(nb_name))}\")\n",
    "    \n",
    "    # script path & name\n",
    "    if folder is not None: folder = Path(folder)\n",
    "    else: folder = nb_path.parent\n",
    "    if name is not None: name = f\"{Path(name).stem}.py\"\n",
    "    else: name = f\"{nb_path.stem}.py\"\n",
    "    script_path = folder/name\n",
    "    \n",
    "    # delete file if exists and create script_path folder if doesn't exist\n",
    "    if os.path.exists(script_path): os.remove(script_path)\n",
    "    script_path.parent.mkdir(parents=True, exist_ok=True)\n",
    "    \n",
    "    # Write script header\n",
    "    with open(script_path, 'w') as f:\n",
    "        f.write(f'# -*- coding: utf-8 -*-\\n')\n",
    "        f.write(f'\"\"\"{nb_name}\\n\\n')\n",
    "        f.write(f'Automatically generated.\\n\\n')\n",
    "        if nb_path is not None:\n",
    "            f.write(f'Original file is located at:\\n')\n",
    "            f.write(f'    {nb_path}\\n')\n",
    "        f.write(f'\"\"\"')\n",
    "\n",
    "    # identify convertible cells (excluding empty and those with hide flags)\n",
    "    for i in range(10):\n",
    "        try: \n",
    "            nb = _read_nb(nb_path)\n",
    "            break\n",
    "        except: \n",
    "            time.sleep(.5)\n",
    "    idxs = _get_unhidden_cells(nb['cells'])\n",
    "    pnb = nbformat.from_dict(nb)\n",
    "    pnb['cells'] = [pnb['cells'][i] for i in idxs]\n",
    "\n",
    "    # clean up cells and write script\n",
    "    sep = '\\n'* 2\n",
    "    for i,cell in enumerate(pnb['cells']):\n",
    "        source_str = cell['source'].replace('\\r', '')\n",
    "        code_lines = source_str.split('\\n')    \n",
    "        if code_lines == ['']: continue\n",
    "        while code_lines[0] == '': code_lines = code_lines[1:]\n",
    "        while code_lines[-1] == '': code_lines = code_lines[:-1]\n",
    "        cl = []\n",
    "        for j in range(len(code_lines)): \n",
    "            if list(set(code_lines[j].split(\" \"))) == ['']:\n",
    "                code_lines[j] = ''\n",
    "            if i == 0 or code_lines[j-1] != '' or code_lines[j] != '': \n",
    "                cl.append(code_lines[j])\n",
    "        code_lines = cl\n",
    "        code = sep + '\\n'.join(code_lines)\n",
    "        with open(script_path, 'a', encoding='utf8') as f: f.write(code)\n",
    "            \n",
    "    # check script exists\n",
    "    assert os.path.isfile(script_path), f\"an error occurred during the export and {script_path} doesn't exist\"\n",
    "    if verbose: \n",
    "        print(f\"{nb_name} converted to {script_path}\")\n",
    "    if run: runpy.run_path(script_path)\n",
    "    return str(script_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/javascript": [
       "IPython.notebook.save_checkpoint();"
      ],
      "text/plain": [
       "<IPython.core.display.Javascript object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "000_export.ipynb saved at 2021-11-29 16:59:05.\n",
      "000_export.ipynb converted to /Users/nacho/notebooks/tsai/nbs/000_export.py\n",
      "000_export.ipynb converted to 000_export.py\n",
      "000_export.ipynb converted to ../nbs/000_export.py\n"
     ]
    },
    {
     "data": {
      "application/javascript": [
       "IPython.notebook.save_checkpoint();"
      ],
      "text/plain": [
       "<IPython.core.display.Javascript object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "000_export.ipynb saved at 2021-11-29 16:59:08.\n",
      "000_export.ipynb converted to ../test_export/000_export.py\n"
     ]
    }
   ],
   "source": [
    "#hide\n",
    "if not is_colab(): \n",
    "\n",
    "    nb = None\n",
    "    folder = None\n",
    "    name = None\n",
    "    pyname = nb2py(nb=nb, folder=folder, name=name)\n",
    "    if pyname is not None: \n",
    "        assert os.path.isfile(pyname)\n",
    "        os.remove(pyname)\n",
    "        assert not os.path.isfile(pyname)\n",
    "\n",
    "    nb = '000_export.ipynb'\n",
    "    folder = None\n",
    "    name = None\n",
    "    pyname = nb2py(nb=nb, folder=folder, name=name)\n",
    "    if pyname is not None: \n",
    "        assert os.path.isfile(pyname)\n",
    "        os.remove(pyname)\n",
    "        assert not os.path.isfile(pyname)\n",
    "\n",
    "    nb = '../nbs/000_export'\n",
    "    folder = None\n",
    "    name = None\n",
    "    pyname = nb2py(nb=nb, folder=folder, name=name)\n",
    "    if pyname is not None: \n",
    "        assert os.path.isfile(pyname)\n",
    "        os.remove(pyname)\n",
    "        assert not os.path.isfile(pyname)\n",
    "\n",
    "    nb = None\n",
    "    folder = '../test_export/'\n",
    "    name = None\n",
    "    pyname = nb2py(nb=nb, folder=folder, name=name)\n",
    "    if pyname is not None: \n",
    "        assert os.path.isfile(pyname)\n",
    "        shutil.rmtree(folder)\n",
    "        assert not os.path.isfile(pyname)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#hide\n",
    "from tsai.imports import create_scripts\n",
    "# from tsai.export import get_nb_name\n",
    "nb_name = get_nb_name()\n",
    "create_scripts(nb_name);"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
